package com.zetatech.bdp.writer.rdbms;

import com.google.common.base.Joiner;
import com.zetatech.bdp.entity.ColumnDef;
import com.zetatech.bdp.entity.RdbmsWriter;
import com.zetatech.bdp.entity.Writer;
import com.zetatech.bdp.writer.WriterHandler;
import com.zetatech.bdp.writer.utils.DataSourceHolder;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author cheng.cao@zetatech.com.cn
 */
@Slf4j
public class RdbmsWriterHandler implements WriterHandler {

    @Override
    public void write(Writer w, List<ColumnDef> columnDefs, List<Map<String, String>> datas) {
        System.out.println("-----------------------" + datas.size());
        RdbmsWriter writer = (RdbmsWriter) w;
        List<String> columnSet = new ArrayList<>();
        for (Map<String, String> data : datas) {
            for (String key : data.keySet()) {
                if (!columnSet.contains(key)) {
                    columnSet.add(key);
                }
            }
        }
        Map<String, Integer> columnIndexMapping = new HashMap<>();
        int index = 1;
        for (String key : columnSet) {
            columnIndexMapping.put(key, index++);
        }
        String questionMark = StringUtils.repeat("?,", columnSet.size());
        questionMark = questionMark.substring(0, questionMark.length() - 1);
        StringBuilder sqlBuilder = new StringBuilder();
        sqlBuilder.append("insert into ").append(writer.getTable());
        sqlBuilder.append("(").append(Joiner.on(",").join(columnSet)).append(")");
        sqlBuilder.append(" values(").append(questionMark).append(")");
        DataSource ds = DataSourceHolder.me().getDataSource(writer.getUrl(), writer.getDriver(), writer.getUsername(), writer.getPassword());
        Connection connection = null;
        try {
            connection = ds.getConnection();
            connection.setAutoCommit(false);
            PreparedStatement pstmt = connection.prepareStatement(sqlBuilder.toString());
            int i = 0;
            for (Map<String, String> line : datas) {
                for (String column : columnSet) {
                    pstmt.setString(columnIndexMapping.get(column), line.get(column));
                }
                pstmt.addBatch();
                if (++i % 5000 == 0) {
                    pstmt.executeBatch();
                }
            }
            if (datas.size() % 5000 != 0) {
                pstmt.executeBatch();
            }
            connection.commit();
        } catch (Exception e) {
            e.printStackTrace();
            try {
                connection.rollback();
            } catch (SQLException ex) {
                ex.printStackTrace();
            }
        } finally {
            try {
                connection.close();
            } catch (Exception e) {
            }
        }
    }
}